package gibbs

import (
	"cs466/findmotif/fasta"
	"cs466/findmotif/motif"
	"fmt"
	"log"
	"math/rand"
	"os"
)

func ExtractAlignment(seqs [][]byte, sites []int, ml int) [][]byte {
	if len(seqs) != len(sites) {
		log.Panicf("wtf? %d, %d", len(seqs), len(sites))
	}
	ret := make([][]byte, len(seqs))

	for i, start := range sites {
		ret[i] = seqs[i][start : start+ml]
	}

	return ret
}

func product(x []float64) (prod float64) {
	prod = 1
	for _, v := range x {
		prod *= v
	}
	return
}

func GibbsSample(seqs [][]byte, ml int) (motif.PWM, []int) {
	SeqLength := len(seqs[0])
	sites := make([]int, len(seqs))
	sc := len(seqs)


	best_ic := 0.
	var best_sites = sites

	const threshold = 2

	// reset and start over sc times; this is arbitrary, but it seems to find
	// the answer almost every time, without running forever
	for n := 0; n < 8000/sc; n++ {
		for i := range sites {
			sites[i] = int(rand.Int31n(int32(SeqLength - ml)))
		}
		for changed := 1+sc; changed > threshold; {
			changed = 0
			for s := range sites {
				sample := make([][]byte, len(seqs)-1)
				sample_sites := make([]int, len(seqs)-1)

				for i := range sample {
					if i < s {
						sample[i] = seqs[i]
						sample_sites[i] = sites[i]
					} else {
						sample[i] = seqs[i+1]
						sample_sites[i] = sites[i+1]
					}
				}

				alignment := ExtractAlignment(sample, sample_sites, ml)
				pwm := PWMFor(alignment)
				probs := CandidateProbabilities(pwm, seqs[s])
				newsite := WeightedSelect(probs)
				if sites[s] != newsite {
					changed++
				}
				sites[s] = newsite

				ic := PWMFor(ExtractAlignment(seqs, sites, ml)).ICPC()
				if ic > best_ic {
					best_ic = ic
					best_sites = make([]int, len(sites))
					copy(best_sites, sites)
				}
			}
		}
	}

	return PWMFor(ExtractAlignment(seqs, best_sites, ml)), best_sites
}

func WritePredictions(dirname string, seqs [][]byte, ml uint) error {
	// perform Gibb's Sampling
	predMotif, predSites := GibbsSample(seqs, int(ml))

	// record predicted motif
	motif_name := fmt.Sprintf("%s/predictedmotif.txt", dirname)
	f, err := os.Create(motif_name)
	if err != nil {
		return err
	}
	defer f.Close()
	err = fasta.WriteMotif(f, predMotif, ml, float64(len(seqs)))
	if err != nil {
		return err
	}

	// record predicted sites
	sites_name := fmt.Sprintf("%s/predictedsites.txt", dirname)
	f, err = os.Create(sites_name)
	if err != nil {
		return err
	}
	defer f.Close()
	_, err = fmt.Fprint(f, predSites)
	if err != nil {
		return err
	}

	return nil
}

func Main(dirname string) {
	// open data files
	seq_name := fmt.Sprintf("%s/sequences.fa", dirname)
	ml_name := fmt.Sprintf("%s/motiflength.txt", dirname)
	seqs, _, err := fasta.LoadFasta(seq_name)
	if err != nil {
		log.Fatal(err)
	}
	ml, err := fasta.LoadMotifLength(ml_name)
	if err != nil {
		log.Fatal(err)
	}

	err = WritePredictions(dirname, seqs, ml)
	if err != nil {
		log.Fatal(err)
	}
}
